#!/usr/bin/env python3
"""
RTSP和HTTP视频流服务器模块

实现标准RTSP协议和HTTP MJPEG流服务器，支持流质量自适应调整。
实现需求1.4, 2.1, 2.3：视频流传输和处理。
"""

import logging
import asyncio
import threading
import time
import socket
import struct
from typing import Optional, Dict, Any, List, Callable, Tuple
from dataclasses import dataclass, field
from enum import Enum
import numpy as np
from datetime import datetime
import io
import base64
import json

try:
    import cv2
    CV2_AVAILABLE = True
except ImportError:
    CV2_AVAILABLE = False
    logging.warning("OpenCV not available, some features may be limited")

from .esp32_camera_driver import ESP32CameraDriver, FrameInfo, CameraState

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class StreamProtocol(Enum):
    """流协议类型"""
    RTSP = "rtsp"
    HTTP = "http"
    MJPEG = "mjpeg"
    WEBSOCKET = "websocket"


class StreamQuality(Enum):
    """流质量等级"""
    LOW = "low"
    MEDIUM = "medium"
    HIGH = "high"
    ULTRA = "ultra"


@dataclass
class StreamConfig:
    """流配置"""
    protocol: StreamProtocol
    port: int
    path: str = "/"
    quality: StreamQuality = StreamQuality.MEDIUM
    max_clients: int = 10
    buffer_size: int = 1024 * 1024  # 1MB
    adaptive_quality: bool = True
    fps_limit: Optional[int] = None
    
    # RTSP特定配置
    rtsp_transport: str = "udp"  # udp, tcp, multicast
    rtsp_session_timeout: int = 60
    
    # HTTP特定配置
    http_boundary: str = "frame"
    http_content_type: str = "multipart/x-mixed-replace"
    
    # 质量自适应配置
    quality_adaptation: Dict[str, Any] = field(default_factory=lambda: {
        'enable': True,
        'min_quality': 10,
        'max_quality': 95,
        'adaptation_interval': 5.0,
        'bandwidth_threshold': 0.8,
        'client_buffer_threshold': 0.7
    })


@dataclass
class ClientInfo:
    """客户端信息"""
    client_id: str
    address: Tuple[str, int]
    protocol: StreamProtocol
    connected_time: datetime
    bytes_sent: int = 0
    frames_sent: int = 0
    last_frame_time: Optional[datetime] = None
    buffer_usage: float = 0.0
    connection_quality: float = 1.0


@dataclass
class StreamStats:
    """流统计信息"""
    active_clients: int = 0
    total_clients_served: int = 0
    bytes_streamed: int = 0
    frames_streamed: int = 0
    current_fps: float = 0.0
    average_bitrate: float = 0.0
    dropped_frames: int = 0
    error_count: int = 0
    uptime: float = 0.0
    last_error: Optional[str] = None


class RTSPServer:
    """RTSP流服务器"""
    
    def __init__(self, config: StreamConfig):
        self.config = config
        self.server_socket: Optional[socket.socket] = None
        self.clients: Dict[str, ClientInfo] = {}
        self.stats = StreamStats()
        self.running = False
        self.server_thread: Optional[threading.Thread] = None
        self._frame_queue = asyncio.Queue(maxsize=10)
        
    async def start(self) -> bool:
        """启动RTSP服务器"""
        try:
            self.server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            self.server_socket.bind(('0.0.0.0', self.config.port))
            self.server_socket.listen(self.config.max_clients)
            
            self.running = True
            self.server_thread = threading.Thread(
                target=self._server_loop,
                name=f"RTSPServer-{self.config.port}",
                daemon=True
            )
            self.server_thread.start()
            
            logger.info(f"RTSP服务器启动成功，端口: {self.config.port}")
            return True
            
        except Exception as e:
            logger.error(f"RTSP服务器启动失败: {e}")
            return False
    
    def _server_loop(self):
        """服务器主循环"""
        while self.running:
            try:
                if self.server_socket:
                    client_socket, address = self.server_socket.accept()
                    client_thread = threading.Thread(
                        target=self._handle_client,
                        args=(client_socket, address),
                        daemon=True
                    )
                    client_thread.start()
            except Exception as e:
                if self.running:
                    logger.error(f"RTSP服务器循环错误: {e}")
                    self.stats.error_count += 1
                    self.stats.last_error = str(e)
    
    def _handle_client(self, client_socket: socket.socket, address: Tuple[str, int]):
        """处理客户端连接"""
        client_id = f"{address[0]}:{address[1]}"
        client_info = ClientInfo(
            client_id=client_id,
            address=address,
            protocol=StreamProtocol.RTSP,
            connected_time=datetime.now()
        )
        
        self.clients[client_id] = client_info
        self.stats.active_clients += 1
        self.stats.total_clients_served += 1
        
        logger.info(f"RTSP客户端连接: {client_id}")
        
        try:
            # 处理RTSP握手
            self._handle_rtsp_handshake(client_socket, client_info)
            
            # 开始流传输
            self._stream_to_client(client_socket, client_info)
            
        except Exception as e:
            logger.error(f"处理RTSP客户端 {client_id} 时发生错误: {e}")
        finally:
            client_socket.close()
            if client_id in self.clients:
                del self.clients[client_id]
            self.stats.active_clients -= 1
            logger.info(f"RTSP客户端断开: {client_id}")
    
    def _handle_rtsp_handshake(self, client_socket: socket.socket, client_info: ClientInfo):
        """处理RTSP握手"""
        # 简化的RTSP握手实现
        # 在实际实现中需要完整的RTSP协议支持
        
        # 接收DESCRIBE请求
        request = client_socket.recv(1024).decode('utf-8')
        logger.debug(f"RTSP请求: {request}")
        
        # 发送SDP响应
        sdp_content = self._generate_sdp()
        response = (
            "RTSP/1.0 200 OK\r\n"
            "CSeq: 1\r\n"
            "Content-Type: application/sdp\r\n"
            f"Content-Length: {len(sdp_content)}\r\n"
            "\r\n"
            f"{sdp_content}"
        )
        client_socket.send(response.encode('utf-8'))
        
        # 处理SETUP请求
        setup_request = client_socket.recv(1024).decode('utf-8')
        setup_response = (
            "RTSP/1.0 200 OK\r\n"
            "CSeq: 2\r\n"
            f"Session: {client_info.client_id}\r\n"
            "Transport: RTP/AVP;unicast;client_port=8000-8001\r\n"
            "\r\n"
        )
        client_socket.send(setup_response.encode('utf-8'))
        
        # 处理PLAY请求
        play_request = client_socket.recv(1024).decode('utf-8')
        play_response = (
            "RTSP/1.0 200 OK\r\n"
            "CSeq: 3\r\n"
            f"Session: {client_info.client_id}\r\n"
            "\r\n"
        )
        client_socket.send(play_response.encode('utf-8'))
    
    def _generate_sdp(self) -> str:
        """生成SDP描述"""
        return (
            "v=0\r\n"
            "o=- 0 0 IN IP4 127.0.0.1\r\n"
            "s=ESP32 Camera Stream\r\n"
            "c=IN IP4 0.0.0.0\r\n"
            "t=0 0\r\n"
            "m=video 0 RTP/AVP 96\r\n"
            "a=rtpmap:96 H264/90000\r\n"
            "a=control:track1\r\n"
        )
    
    def _stream_to_client(self, client_socket: socket.socket, client_info: ClientInfo):
        """向客户端流式传输数据"""
        # 这里应该实现RTP数据包的发送
        # 简化实现，实际需要完整的RTP/RTCP协议支持
        logger.info(f"开始向客户端 {client_info.client_id} 流式传输")
        
        while self.running:
            try:
                # 模拟发送RTP数据包
                time.sleep(1.0 / 30)  # 30fps
                client_info.frames_sent += 1
                client_info.last_frame_time = datetime.now()
                
            except Exception as e:
                logger.error(f"流式传输到客户端 {client_info.client_id} 失败: {e}")
                break
    
    async def add_frame(self, frame: np.ndarray, frame_info: FrameInfo):
        """添加帧到流队列"""
        try:
            if not self._frame_queue.full():
                await self._frame_queue.put((frame, frame_info))
            else:
                self.stats.dropped_frames += 1
        except Exception as e:
            logger.error(f"添加帧到RTSP队列失败: {e}")
    
    def stop(self):
        """停止RTSP服务器"""
        self.running = False
        if self.server_socket:
            self.server_socket.close()
        logger.info("RTSP服务器已停止")


class HTTPStreamServer:
    """HTTP MJPEG流服务器"""
    
    def __init__(self, config: StreamConfig):
        self.config = config
        self.clients: Dict[str, ClientInfo] = {}
        self.stats = StreamStats()
        self.running = False
        self.app = None
        self._frame_queue = asyncio.Queue(maxsize=10)
        self._current_frame: Optional[bytes] = None
        self._frame_lock = threading.Lock()
        
    async def start(self) -> bool:
        """启动HTTP流服务器"""
        try:
            from aiohttp import web, WSMsgType
            from aiohttp.web_response import StreamResponse
            
            self.app = web.Application()
            self.app.router.add_get(self.config.path, self._handle_stream_request)
            self.app.router.add_get('/status', self._handle_status_request)
            self.app.router.add_get('/snapshot', self._handle_snapshot_request)
            
            runner = web.AppRunner(self.app)
            await runner.setup()
            
            site = web.TCPSite(runner, '0.0.0.0', self.config.port)
            await site.start()
            
            self.running = True
            logger.info(f"HTTP流服务器启动成功，端口: {self.config.port}")
            return True
            
        except Exception as e:
            logger.error(f"HTTP流服务器启动失败: {e}")
            return False
    
    async def _handle_stream_request(self, request):
        """处理流请求"""
        from aiohttp.web_response import StreamResponse
        
        client_id = f"{request.remote}:{request.transport.get_extra_info('peername')[1] if request.transport else 'unknown'}"
        client_info = ClientInfo(
            client_id=client_id,
            address=(request.remote, 0),
            protocol=StreamProtocol.HTTP,
            connected_time=datetime.now()
        )
        
        self.clients[client_id] = client_info
        self.stats.active_clients += 1
        self.stats.total_clients_served += 1
        
        logger.info(f"HTTP客户端连接: {client_id}")
        
        # 创建流响应
        response = StreamResponse(
            status=200,
            reason='OK',
            headers={
                'Content-Type': f'{self.config.http_content_type}; boundary={self.config.http_boundary}',
                'Cache-Control': 'no-cache',
                'Connection': 'close',
                'Pragma': 'no-cache'
            }
        )
        
        await response.prepare(request)
        
        try:
            while self.running:
                # 获取当前帧
                frame_data = self._get_current_frame()
                if frame_data:
                    # 发送MJPEG帧
                    boundary_data = f"\r\n--{self.config.http_boundary}\r\n".encode()
                    header_data = f"Content-Type: image/jpeg\r\nContent-Length: {len(frame_data)}\r\n\r\n".encode()
                    
                    await response.write(boundary_data)
                    await response.write(header_data)
                    await response.write(frame_data)
                    
                    client_info.frames_sent += 1
                    client_info.bytes_sent += len(frame_data)
                    client_info.last_frame_time = datetime.now()
                
                # 控制帧率
                fps_limit = self.config.fps_limit or 30
                await asyncio.sleep(1.0 / fps_limit)
                
        except Exception as e:
            logger.error(f"HTTP流传输到客户端 {client_id} 失败: {e}")
        finally:
            if client_id in self.clients:
                del self.clients[client_id]
            self.stats.active_clients -= 1
            logger.info(f"HTTP客户端断开: {client_id}")
        
        return response
    
    async def _handle_status_request(self, request):
        """处理状态请求"""
        from aiohttp import web
        
        status = {
            'server': 'HTTP MJPEG Stream Server',
            'running': self.running,
            'config': {
                'port': self.config.port,
                'path': self.config.path,
                'quality': self.config.quality.value,
                'max_clients': self.config.max_clients
            },
            'stats': {
                'active_clients': self.stats.active_clients,
                'total_clients_served': self.stats.total_clients_served,
                'bytes_streamed': self.stats.bytes_streamed,
                'frames_streamed': self.stats.frames_streamed,
                'dropped_frames': self.stats.dropped_frames,
                'error_count': self.stats.error_count
            },
            'clients': [
                {
                    'id': client.client_id,
                    'address': f"{client.address[0]}:{client.address[1]}",
                    'connected_time': client.connected_time.isoformat(),
                    'frames_sent': client.frames_sent,
                    'bytes_sent': client.bytes_sent
                }
                for client in self.clients.values()
            ]
        }
        
        return web.json_response(status)
    
    async def _handle_snapshot_request(self, request):
        """处理快照请求"""
        from aiohttp import web
        
        frame_data = self._get_current_frame()
        if frame_data:
            return web.Response(
                body=frame_data,
                content_type='image/jpeg',
                headers={'Cache-Control': 'no-cache'}
            )
        else:
            return web.Response(
                text='No frame available',
                status=503,
                content_type='text/plain'
            )
    
    def _get_current_frame(self) -> Optional[bytes]:
        """获取当前帧数据"""
        with self._frame_lock:
            return self._current_frame
    
    async def add_frame(self, frame: np.ndarray, frame_info: FrameInfo):
        """添加帧到流"""
        try:
            # 将帧编码为JPEG
            if CV2_AVAILABLE:
                # 使用OpenCV编码
                encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), self._get_jpeg_quality()]
                _, buffer = cv2.imencode('.jpg', frame, encode_param)
                frame_data = buffer.tobytes()
            else:
                # 简化编码（实际应用中需要更好的编码方案）
                frame_data = self._simple_jpeg_encode(frame)
            
            # 更新当前帧
            with self._frame_lock:
                self._current_frame = frame_data
            
            # 更新统计信息
            self.stats.frames_streamed += 1
            self.stats.bytes_streamed += len(frame_data)
            
        except Exception as e:
            logger.error(f"添加帧到HTTP流失败: {e}")
            self.stats.error_count += 1
    
    def _get_jpeg_quality(self) -> int:
        """根据质量等级获取JPEG质量参数"""
        quality_map = {
            StreamQuality.LOW: 30,
            StreamQuality.MEDIUM: 60,
            StreamQuality.HIGH: 85,
            StreamQuality.ULTRA: 95
        }
        return quality_map.get(self.config.quality, 60)
    
    def _simple_jpeg_encode(self, frame: np.ndarray) -> bytes:
        """简单的JPEG编码（当OpenCV不可用时）"""
        # 这是一个非常简化的实现，实际应用中需要真正的JPEG编码
        # 这里只是返回一个占位符
        return b'\xff\xd8\xff\xe0\x00\x10JFIF\x00\x01\x01\x01\x00H\x00H\x00\x00\xff\xdb\x00C\x00\x08\x06\x06\x07\x06\x05\x08\x07\x07\x07\t\t\x08\n\x0c\x14\r\x0c\x0b\x0b\x0c\x19\x12\x13\x0f\x14\x1d\x1a\x1f\x1e\x1d\x1a\x1c\x1c $.\' ",#\x1c\x1c(7),01444\x1f\'9=82<.342\xff\xc0\x00\x11\x08\x00\x01\x00\x01\x01\x01\x11\x00\x02\x11\x01\x03\x11\x01\xff\xc4\x00\x14\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x08\xff\xc4\x00\x14\x10\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\xff\xda\x00\x0c\x03\x01\x00\x02\x11\x03\x11\x00\x3f\x00\xaa\xff\xd9'
    
    def stop(self):
        """停止HTTP服务器"""
        self.running = False
        logger.info("HTTP流服务器已停止")


class StreamQualityManager:
    """流质量管理器"""
    
    def __init__(self, config: StreamConfig):
        self.config = config
        self.current_quality = config.quality
        self.adaptation_enabled = config.quality_adaptation.get('enable', True)
        self.last_adaptation_time = time.time()
        self.quality_history: List[Tuple[float, StreamQuality]] = []
        
    def update_client_metrics(self, clients: Dict[str, ClientInfo]):
        """更新客户端指标"""
        if not self.adaptation_enabled:
            return
        
        current_time = time.time()
        adaptation_interval = self.config.quality_adaptation.get('adaptation_interval', 5.0)
        
        if current_time - self.last_adaptation_time < adaptation_interval:
            return
        
        # 计算平均连接质量
        if clients:
            avg_quality = sum(client.connection_quality for client in clients.values()) / len(clients)
            avg_buffer_usage = sum(client.buffer_usage for client in clients.values()) / len(clients)
            
            # 根据连接质量调整流质量
            new_quality = self._calculate_optimal_quality(avg_quality, avg_buffer_usage)
            
            if new_quality != self.current_quality:
                logger.info(f"自适应质量调整: {self.current_quality.value} -> {new_quality.value}")
                self.current_quality = new_quality
                self.quality_history.append((current_time, new_quality))
                
                # 保持历史记录在合理范围内
                if len(self.quality_history) > 100:
                    self.quality_history = self.quality_history[-50:]
        
        self.last_adaptation_time = current_time
    
    def _calculate_optimal_quality(self, avg_quality: float, avg_buffer_usage: float) -> StreamQuality:
        """计算最优质量等级"""
        bandwidth_threshold = self.config.quality_adaptation.get('bandwidth_threshold', 0.8)
        buffer_threshold = self.config.quality_adaptation.get('client_buffer_threshold', 0.7)
        
        # 如果连接质量或缓冲区使用率过高，降低质量
        if avg_quality < bandwidth_threshold or avg_buffer_usage > buffer_threshold:
            if self.current_quality == StreamQuality.ULTRA:
                return StreamQuality.HIGH
            elif self.current_quality == StreamQuality.HIGH:
                return StreamQuality.MEDIUM
            elif self.current_quality == StreamQuality.MEDIUM:
                return StreamQuality.LOW
        
        # 如果连接质量良好且缓冲区使用率低，提高质量
        elif avg_quality > 0.9 and avg_buffer_usage < 0.3:
            if self.current_quality == StreamQuality.LOW:
                return StreamQuality.MEDIUM
            elif self.current_quality == StreamQuality.MEDIUM:
                return StreamQuality.HIGH
            elif self.current_quality == StreamQuality.HIGH:
                return StreamQuality.ULTRA
        
        return self.current_quality
    
    def get_current_quality(self) -> StreamQuality:
        """获取当前质量等级"""
        return self.current_quality
    
    def set_quality(self, quality: StreamQuality):
        """手动设置质量等级"""
        self.current_quality = quality
        logger.info(f"手动设置流质量: {quality.value}")


class VideoStreamManager:
    """视频流管理器"""
    
    def __init__(self):
        self.servers: Dict[str, Any] = {}
        self.camera_driver: Optional[ESP32CameraDriver] = None
        self.quality_managers: Dict[str, StreamQualityManager] = {}
        self.running = False
        
    def add_rtsp_server(self, name: str, config: StreamConfig) -> bool:
        """添加RTSP服务器"""
        try:
            server = RTSPServer(config)
            quality_manager = StreamQualityManager(config)
            
            self.servers[name] = server
            self.quality_managers[name] = quality_manager
            
            logger.info(f"RTSP服务器 '{name}' 已添加")
            return True
        except Exception as e:
            logger.error(f"添加RTSP服务器 '{name}' 失败: {e}")
            return False
    
    def add_http_server(self, name: str, config: StreamConfig) -> bool:
        """添加HTTP服务器"""
        try:
            server = HTTPStreamServer(config)
            quality_manager = StreamQualityManager(config)
            
            self.servers[name] = server
            self.quality_managers[name] = quality_manager
            
            logger.info(f"HTTP服务器 '{name}' 已添加")
            return True
        except Exception as e:
            logger.error(f"添加HTTP服务器 '{name}' 失败: {e}")
            return False
    
    def set_camera_driver(self, driver: ESP32CameraDriver):
        """设置摄像头驱动"""
        self.camera_driver = driver
        
        # 设置帧回调
        driver.set_frame_callback(self._on_frame_received)
    
    def _on_frame_received(self, frame: np.ndarray, frame_info: FrameInfo):
        """处理接收到的帧"""
        # 将帧分发到所有服务器
        for name, server in self.servers.items():
            try:
                asyncio.create_task(server.add_frame(frame, frame_info))
                
                # 更新质量管理
                if name in self.quality_managers:
                    self.quality_managers[name].update_client_metrics(server.clients)
                    
            except Exception as e:
                logger.error(f"向服务器 '{name}' 分发帧失败: {e}")
    
    async def start_all_servers(self) -> bool:
        """启动所有服务器"""
        success_count = 0
        
        for name, server in self.servers.items():
            try:
                if await server.start():
                    success_count += 1
                    logger.info(f"服务器 '{name}' 启动成功")
                else:
                    logger.error(f"服务器 '{name}' 启动失败")
            except Exception as e:
                logger.error(f"启动服务器 '{name}' 时发生异常: {e}")
        
        self.running = success_count > 0
        return self.running
    
    def stop_all_servers(self):
        """停止所有服务器"""
        for name, server in self.servers.items():
            try:
                server.stop()
                logger.info(f"服务器 '{name}' 已停止")
            except Exception as e:
                logger.error(f"停止服务器 '{name}' 时发生异常: {e}")
        
        self.running = False
    
    def get_server_status(self, name: str) -> Optional[Dict[str, Any]]:
        """获取服务器状态"""
        if name not in self.servers:
            return None
        
        server = self.servers[name]
        quality_manager = self.quality_managers.get(name)
        
        return {
            'name': name,
            'running': getattr(server, 'running', False),
            'stats': server.stats.__dict__ if hasattr(server, 'stats') else {},
            'clients': len(getattr(server, 'clients', {})),
            'current_quality': quality_manager.get_current_quality().value if quality_manager else None
        }
    
    def get_all_status(self) -> Dict[str, Any]:
        """获取所有服务器状态"""
        return {
            'manager_running': self.running,
            'camera_connected': self.camera_driver is not None and self.camera_driver.state == CameraState.STREAMING,
            'servers': {
                name: self.get_server_status(name)
                for name in self.servers.keys()
            }
        }


async def main():
    """测试函数"""
    # 创建流管理器
    stream_manager = VideoStreamManager()
    
    # 添加HTTP服务器
    http_config = StreamConfig(
        protocol=StreamProtocol.HTTP,
        port=8080,
        path="/stream",
        quality=StreamQuality.MEDIUM,
        adaptive_quality=True
    )
    stream_manager.add_http_server("main_http", http_config)
    
    # 添加RTSP服务器
    rtsp_config = StreamConfig(
        protocol=StreamProtocol.RTSP,
        port=8554,
        path="/stream",
        quality=StreamQuality.HIGH
    )
    stream_manager.add_rtsp_server("main_rtsp", rtsp_config)
    
    # 启动所有服务器
    if await stream_manager.start_all_servers():
        print("所有流服务器启动成功")
        
        # 显示状态
        status = stream_manager.get_all_status()
        print(f"流管理器状态: {json.dumps(status, indent=2, ensure_ascii=False)}")
        
        # 运行一段时间
        await asyncio.sleep(10)
        
        # 停止服务器
        stream_manager.stop_all_servers()
        print("所有流服务器已停止")
    else:
        print("流服务器启动失败")


if __name__ == "__main__":
    asyncio.run(main())